import os
import sys
import time
import base64
import pprint
import random
import sqlite3
from hashlib import md5
from xmlrpclib import ServerProxy

import pefile
import peutils
import isexe

MAX_RETRIES = 5
CONFIG_FILENAME = "servers.conf"
TIMEOUT = 10
MIN_STRING = 9

def doLog(msg):
	print "[%s] %s" %(time.ctime(), msg)

def checkPacker(pe):
	try:
		sig = peutils.SignatureDatabase("UserDB.TXT")
		matches = sig.match_all(pe, ep_only = True)
	except:
		matches = None

	if not matches:
		return None

	return matches
	
def bufToFile(buf, filename):
	f = file(filename, "wb")
	f.write(base64.b64decode(buf))
	f.close()
	
def createSchema(conn):
	try:
		sql = """ create table unpacking (md5, filename, packer, correct) """
		cur = conn.cursor()
		cur.execute(sql)
	except:
		pass

def updateDatabase(filename, packers, correct):
	conn = sqlite3.connect("packers.sqlite")
	createSchema(conn)
	
	sql = """ insert into unpacking values (?, ?, ?, ?) """
	cur = conn.cursor()
	values = (md5(file(filename, "rb").read()).hexdigest(), filename, str(packers), str(correct))
	cur.execute(sql, values)
	conn.commit()
	cur.close()
	
	conn.close()

def connect(url):
	if url == "auto": # Not an URL, but a configuration option
		servers = file(CONFIG_FILENAME, "r").readlines()
		# Randomize the list for load balancing
		random.shuffle(servers)
		
		for server in servers:
			if server == "" or server.startswith("#"):
				continue
			
			s = connect(server)
			
			if s is not None:
				return s
		
		doLog("No server is available!")
		return None
	else:
		url = url.strip("\r").strip("\n")
		s = ServerProxy(url, allow_none = True)
		
		try:
			status = s.ping()
			doLog("Checking server %s status: %s" %(url, status))
			
			if status != "[ALIVE]":
				return None
		except:
			doLog("Fatal Error: " + str(sys.exc_info()[1]))
			return None
		
		return s

def dump_info(self):
	"""This is a simplified version of PE.dumpinfo(self)"""
	
	dump = pefile.Dump()
	warnings = self.get_warnings()
	if warnings:
		dump.add_header('Parsing Warnings')
		for warning in warnings:
			dump.add_line(warning)
			
	dump.add_header('PE Sections')
	
	section_flags = pefile.retrieve_flags(pefile.SECTION_CHARACTERISTICS, 'IMAGE_SCN_')
	
	for section in self.sections:
		dump.add_lines(section.dump())
		dump.add('Flags: ')
		flags = []
		for flag in section_flags:
			if getattr(section, flag[0]):
				flags.append(flag[0])
		dump.add_line(', '.join(flags))
		dump.add_line('Entropy: %f (Min=0.0, Max=8.0)' % section.get_entropy() )
		if pefile.md5 is not None:
			dump.add_line('MD5     hash: %s' % section.get_hash_md5() )
		dump.add_newline()
			
	if hasattr(self, 'DIRECTORY_ENTRY_IMPORT'):
		dump.add_header('Imported symbols')
		for module in self.DIRECTORY_ENTRY_IMPORT:
			for symbol in module.imports:
				if symbol.import_by_ordinal is True:
					dump.add('%s Ordinal[%s] (Imported by Ordinal)' % (
						module.dll, str(symbol.ordinal)))
				else:
					dump.add('%s.%s Hint[%s]' % (
						module.dll, symbol.name, str(symbol.hint)))
				
				if symbol.bound:
					dump.add_line(' Bound: 0x%08X' % (symbol.bound))
				else:
					dump.add_newline()
			dump.add_newline()
        
        
	if hasattr(self, 'DIRECTORY_ENTRY_BOUND_IMPORT'):
		dump.add_header('Bound imports')
		for bound_imp_desc in self.DIRECTORY_ENTRY_BOUND_IMPORT:
			dump.add_line('DLL: %s' % bound_imp_desc.name)
			dump.add_newline()
			
			for bound_imp_ref in bound_imp_desc.entries:
				dump.add_lines(bound_imp_ref.struct.dump(), 4)
				dump.add_line('DLL: %s' % bound_imp_ref.name, 4)
				dump.add_newline()

	
	if hasattr(self, 'DIRECTORY_ENTRY_DELAY_IMPORT'):
		dump.add_header('Delay Imported symbols')
		for module in self.DIRECTORY_ENTRY_DELAY_IMPORT:
			for symbol in module.imports:
				if symbol.import_by_ordinal is True:
					dump.add('%s Ordinal[%s] (Imported by Ordinal)' % (
						module.dll, str(symbol.ordinal)))
				else:
					dump.add('%s.%s Hint[%s]' % (
						module.dll, symbol.name, str(symbol.hint)))
				
				if symbol.bound:
					dump.add_line(' Bound: 0x%08X' % (symbol.bound))
				else:
					dump.add_newline()
			dump.add_newline()
	
	return dump.get_text()
		
def staticAnalyze(filename, output):
	"""Extract static report and packers from PE file into the output directory"""

	try:
		pe = pefile.PE(filename)
	except pefile.PEFormatError:
		doLog("PEFormatError in static analysis")
		return None
		
	with open(output + "report_static.txt", "w") as out:
		# out.write(pe.dump_info())
		out.write(dump_info(pe))
	
	doLog("Writting static report [DONE]")
	
	packers = checkPacker(pe)
	if packers:
		basename = os.path.basename(filename)
		doLog("Dumping file %s packed with:" % basename)
		pprint.pprint(packers)
		with open(output + "report_packers.txt", "w") as out:
			pprint.pprint(packers, out)
		doLog("Writting packers report [DONE]")
	return packers
		
def main(url, filename, output, retries = 0):
	if retries > MAX_RETRIES:
		doLog("Error: Too many retries (%d), exitting..." % MAX_RETRIES)
	
	if not isexe.main(filename):
		return
	
	buf = file(filename, "rb").read()
	origbuf = buf
	buf = base64.b64encode(buf)
	
	name_id = str(md5(origbuf).hexdigest()) + "_" + str(time.time())
	output = output + os.sep + name_id + os.sep
	os.mkdir(output)
	
	packers = staticAnalyze(filename, output)
	
	s = connect(url)

	if s is None:
		doLog("No server available")
		return
	
	with open(output + "strings_filtered.txt", "w") as out_file:
		strings_file = output + "strings.txt"
		os.system("tools\strings.exe -n %d -q %s > %s" % (MIN_STRING, filename, strings_file))
		with open(strings_file) as in_file:
			line_list = []
			for line in in_file:
				if len(line) > MIN_STRING and not line in line_list:
					line_list.append(line)
					out_file.write(line)
	
	try:
		ret = s.dump(buf, TIMEOUT)
	except:
		doLog("Error running dump (%s), retry number %d..." % (sys.exc_info()[1], retries))
		return
		# main(url, filename, output, retries+1)
	
	i = 0
	if ret is None:
		doLog("Error: No response received!")
		return
		
	if ret["report_warning"]:
		i += 1
		with open(output + "report_warning.txt", "w") as out:
			out.write(ret["report_warning"].encode("UTF-8"))
		doLog("Writting warning report [DONE]")
		
	if ret["report"]:
		i += 1
		report_name = output + "report.txt"
		bufToFile(ret["report"], report_name)
		doLog("Writting behavior report [DONE]")
		
	if ret["report_filtered"]:
		i += 1
		report_name = output + "report_filtered.txt"
		bufToFile(ret["report_filtered"], report_name)
		doLog("Writting behavior filtered report [DONE]")
		
	if ret["report_signature"]:
		i += 1
		report_name = output + "report_signature.txt"
		bufToFile(ret["report_signature"], report_name)
		doLog("Writting behavior signature report [DONE]")
		
	if ret["dump"]:
		i += 1
		report_name = output + "malware.dmp"
		bufToFile(ret["dump"], report_name)
		doLog("Writting dump file [DONE]")
		
	if ret["tricks_debug"]:
		i += 1
		with open(output + "tricks_debug.txt", "w") as out:
			for each in ret["tricks_debug"]:
				out.write(each + "\n")
		doLog("Writting debugging tricks file [DONE]")
	
	if ret["tricks_vm"]:
		i += 1
		with open(output + "tricks_vm.txt", "w") as out:
			for each in ret["tricks_vm"]:
				out.write(each + "\n")
		doLog("Writting VM detection tricks file [DONE]")
		
	if i == 0:
		doLog("Dumping: [FAILED]")
		updateDatabase(filename, packers, False)
	else:
		doLog("Dumping: [DONE]")
		updateDatabase(filename, packers, True)

def usage():
	print "Usage:", sys.argv[0], "<url | auto> <filename> <output directory>"
	print

if __name__ == "__main__":
	if len(sys.argv) < 4:
		usage()
		sys.exit(0)
	else:
		main(sys.argv[1], sys.argv[2], sys.argv[3])
